{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d047044f",
   "metadata": {},
   "source": [
    "# Agents with Memory\n",
    "\n",
    "We have an email assistant that uses a router to triage emails and then passes the email to the agent for response generation. We've also evaluated it and added human-in-the-loop (HITL) to review specific tool calls. Now, we add memory, giving our assistant the ability to remember our HITL feedback!\n",
    "\n",
    "![overview-img](img/overview_memory.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "143094b4",
   "metadata": {},
   "source": [
    "#### Load Environment Variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "085c21ad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 77,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dotenv import load_dotenv\n",
    "load_dotenv(\"../.env\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b99855af",
   "metadata": {},
   "source": [
    "## Memory in LangGraph\n",
    "\n",
    "### Thread-Scoped and Across-Thread Memory\n",
    "\n",
    "First, it's worth explaining how [memory works in LangGraph](https://langchain-ai.github.io/langgraph/concepts/memory/). LangGraph offers two distinct types of memory that serve complementary purposes:\n",
    "\n",
    "**Thread-Scoped Memory (Short-term)** operates within the boundaries of a single conversation thread. It's automatically managed as part of the graph's state and persisted through thread-scoped checkpoints. This memory type retains conversation history, uploaded files, retrieved documents, and other artifacts generated during the interaction. Think of it as the working memory that maintains context within one specific conversation, allowing the agent to reference earlier messages or actions without starting from scratch each time.\n",
    "\n",
    "**Across-Thread Memory (Long-term)** extends beyond individual conversations, creating a persistent knowledge base that spans multiple sessions. This memory is stored as JSON documents in a memory store, organized by namespaces (like folders) and distinct keys (like filenames). Unlike thread-scoped memory, this information persists even after conversations end, enabling the system to recall user preferences, past decisions, and accumulated knowledge. This is what allows an agent to truly learn and adapt over time, rather than treating each interaction as isolated.\n",
    "\n",
    "![short-vs-long-term-memory](img/short-vs-long.png)\n",
    "\n",
    "The [Store](https://langchain-ai.github.io/langgraph/reference/store/#langgraph.store.base.BaseStore) is the foundation of this architecture, providing a flexible database where memories can be organized, retrieved, and updated. What makes this approach powerful is that regardless of which memory type you're working with, the same Store interface provides consistent access patterns. This allows your agent's code to remain unchanged whether you're using a simple in-memory implementation during development or a production-grade database in deployment. \n",
    "\n",
    "### LangGraph Store\n",
    "\n",
    "LangGraph offers different [Store implementations](https://langchain-ai.github.io/langgraph/reference/store/#langgraph.store.base.BaseStore) depending on your [deployment](https://langchain-ai.github.io/langgraph/tutorials/deployment/#other-deployment-options):\n",
    "\n",
    "1. **In-Memory (e.g., notebooks)**:\n",
    "   - Uses `from langgraph.store.memory import InMemoryStore`\n",
    "   - Purely a Python dictionary in memory with no persistence\n",
    "   - Data is lost when the process terminates\n",
    "   - Useful for quick experiments and testing\n",
    "   - Semantic search can be configured as shown [here](https://langchain-ai.github.io/langgraph/how-tos/memory/semantic-search/)\n",
    "\n",
    "2. **Local Development with `langgraph dev`**:\n",
    "   - Similar to InMemoryStore but with pseudo-persistence\n",
    "   - Data is pickled to the local filesystem between restarts\n",
    "   - Lightweight and fast, no need for external databases\n",
    "   - Semantic search can be configured as shown [here](https://langchain-ai.github.io/langgraph/cloud/deployment/semantic_search/)\n",
    "   - Great for development but not designed for production use\n",
    "\n",
    "3. **LangGraph Platform or Production Deployments**:\n",
    "   - Uses PostgreSQL with pgvector for production-grade persistence\n",
    "   - Fully persistent data storage with reliable backups\n",
    "   - Scalable for larger datasets\n",
    "   - Semantic search can be configured as shown [here](https://langchain-ai.github.io/langgraph/cloud/deployment/semantic_search/)\n",
    "   - Default distance metric is cosine similarity (customizable)\n",
    "\n",
    "Let's use the `InMemoryStore` here in the notebook! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "7fa1dda7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.store.memory import InMemoryStore\n",
    "in_memory_store = InMemoryStore()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aceb204c",
   "metadata": {},
   "source": [
    "Memories are namespaced by a tuple, which in this specific example will be (`<user_id>`, \"memories\"). The namespace can be any length and represent anything, it does not have to be user specific."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "f0488a5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "user_id = \"1\"\n",
    "namespace_for_memory = (user_id, \"memories\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3da8b303",
   "metadata": {},
   "source": [
    "We use the `store.put` method to save memories to our namespace in the store. When we do this, we specify the namespace, as defined above, and a key-value pair for the memory: the key is simply a unique identifier for the memory (memory_id) and the value (a dictionary) is the memory itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "4af95b32",
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "memory_id = str(uuid.uuid4())\n",
    "memory = {\"food_preference\" : \"I like pizza\"}\n",
    "in_memory_store.put(namespace_for_memory, memory_id, memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60408492",
   "metadata": {},
   "source": [
    "We can read out memories in our namespace using the `store.search` method, which will return all memories for a given user as a list. The most recent memory is the last in the list. Each memory type is a Python class (`Item`) with certain attributes. We can access it as a dictionary by converting via `.dict`. The attributes it has are shown below, but the most important one is typically `value`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "4c25f5ef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'namespace': ['1', 'memories'],\n",
       " 'key': 'afeb46b1-9a2c-4b10-b52d-e53f93f7fb7a',\n",
       " 'value': {'food_preference': 'I like pizza'},\n",
       " 'created_at': '2025-06-05T18:21:36.440950+00:00',\n",
       " 'updated_at': '2025-06-05T18:21:36.440955+00:00',\n",
       " 'score': None}"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memories = in_memory_store.search(namespace_for_memory)\n",
    "memories[-1].dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44f3e781",
   "metadata": {},
   "source": [
    "To use this in a graph, all we need to do is compile the graph with the store:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "6476b361",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We need this because we want to enable threads (conversations)\n",
    "from langgraph.checkpoint.memory import InMemorySaver\n",
    "checkpointer = InMemorySaver()\n",
    "# We need this because we want to enable across-thread memory\n",
    "from langgraph.store.memory import InMemoryStore\n",
    "in_memory_store = InMemoryStore()\n",
    "# Compile the graph with the checkpointer and store\n",
    "# graph = graph.compile(checkpointer=checkpointer, store=in_memory_store)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c982928",
   "metadata": {},
   "source": [
    "The store is then accessible in any node of the graph, as we'll see below!\n",
    "\n",
    "## Adding Memory to our Assistant\n",
    "\n",
    "Let's take our graph with HITL and add memory to it. This will be very similar to what we had previously. We'll simply update memory in the store when we get feedback from the user.\n",
    "\n",
    "![overview-img](img/HITL_flow_memory.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "38308fc3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv(\".env\", override=True)\n",
    "\n",
    "from typing import Literal\n",
    "from datetime import datetime\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "from langchain.chat_models import init_chat_model\n",
    "from langchain_core.tools import tool\n",
    "\n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.store.base import BaseStore\n",
    "from langgraph.types import interrupt, Command\n",
    "\n",
    "from email_assistant.prompts import triage_system_prompt, triage_user_prompt, agent_system_prompt_hitl_memory, default_triage_instructions, default_background, default_response_preferences, default_cal_preferences, MEMORY_UPDATE_INSTRUCTIONS, MEMORY_UPDATE_INSTRUCTIONS_REINFORCEMENT\n",
    "from email_assistant.tools.default.prompt_templates import HITL_MEMORY_TOOLS_PROMPT\n",
    "from email_assistant.schemas import State, RouterSchema, StateInput\n",
    "from email_assistant.utils import parse_email, format_for_display, format_email_markdown\n",
    "\n",
    "# Agent tools \n",
    "@tool\n",
    "def write_email(to: str, subject: str, content: str) -> str:\n",
    "    \"\"\"Write and send an email.\"\"\"\n",
    "    # Placeholder response - in real app would send email\n",
    "    return f\"Email sent to {to} with subject '{subject}' and content: {content}\"\n",
    "\n",
    "@tool\n",
    "def schedule_meeting(\n",
    "    attendees: list[str], subject: str, duration_minutes: int, preferred_day: datetime, start_time: int\n",
    ") -> str:\n",
    "    \"\"\"Schedule a calendar meeting.\"\"\"\n",
    "    # Placeholder response - in real app would check calendar and schedule\n",
    "    date_str = preferred_day.strftime(\"%A, %B %d, %Y\")\n",
    "    return f\"Meeting '{subject}' scheduled on {date_str} at {start_time} for {duration_minutes} minutes with {len(attendees)} attendees\"\n",
    "\n",
    "@tool\n",
    "def check_calendar_availability(day: str) -> str:\n",
    "    \"\"\"Check calendar availability for a given day.\"\"\"\n",
    "    # Placeholder response - in real app would check actual calendar\n",
    "    return f\"Available times on {day}: 9:00 AM, 2:00 PM, 4:00 PM\"\n",
    "\n",
    "@tool\n",
    "class Question(BaseModel):\n",
    "      \"\"\"Question to ask user.\"\"\"\n",
    "      content: str\n",
    "\n",
    "@tool\n",
    "class Done(BaseModel):\n",
    "      \"\"\"E-mail has been sent.\"\"\"\n",
    "      done: bool\n",
    "    \n",
    "# All tools available to the agent\n",
    "tools = [\n",
    "    write_email, \n",
    "    schedule_meeting, \n",
    "    check_calendar_availability, \n",
    "    Question, \n",
    "    Done\n",
    "]\n",
    "\n",
    "tools_by_name = {tool.name: tool for tool in tools}\n",
    "\n",
    "# Initialize the LLM for use with router / structured output\n",
    "llm = init_chat_model(\"openai:gpt-4.1\", temperature=0.0)\n",
    "llm_router = llm.with_structured_output(RouterSchema) \n",
    "\n",
    "# Initialize the LLM, enforcing tool use (of any available tools) for agent\n",
    "llm = init_chat_model(\"openai:gpt-4.1\", temperature=0.0)\n",
    "llm_with_tools = llm.bind_tools(tools, tool_choice=\"required\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03538f56",
   "metadata": {},
   "source": [
    "Now, this is the critical part! We currently don't capture any feedback from the user in our graph. \n",
    "\n",
    "### Memory Management \n",
    "\n",
    "What we *want* to do is fairly straightforward: we want to add the feedback to the memory `Store`. If we compile our graph with the `Store`, we can access it in any node. So that is not a problem! \n",
    "\n",
    "But we have to answer two questions: \n",
    "\n",
    "1) how do we want the memory to be structured?\n",
    "2) how do we want to update the memory?\n",
    "\n",
    "For 1) we'll just store memories as string to keep things simple. In the below function, we'll just fetch memories from the store as string and initialize with default if it doesn't exist."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "d2715152-2d19-4449-be4b-fdc602eee52d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"\\nEmails that are not worth responding to:\\n- Marketing newsletters and promotional emails\\n- Spam or suspicious emails\\n- CC'd on FYI threads with no direct questions\\n\\nThere are also other things that should be known about, but don't require an email response. For these, you should notify (using the `notify` response). Examples of this include:\\n- Team member out sick or on vacation\\n- Build system notifications or deployments\\n- Project status updates without action items\\n- Important company announcements\\n- FYI emails that contain relevant information for current projects\\n- HR Department deadline reminders\\n- Subscription status / renewal reminders\\n- GitHub notifications\\n\\nEmails that are worth responding to:\\n- Direct questions from team members requiring expertise\\n- Meeting requests requiring confirmation\\n- Critical bug reports related to team's projects\\n- Requests from management requiring acknowledgment\\n- Client inquiries about project status or features\\n- Technical questions about documentation, code, or APIs (especially questions about missing endpoints or features)\\n- Personal reminders related to family (wife / daughter)\\n- Personal reminder related to self-care (doctor appointments, etc)\\n\""
      ]
     },
     "execution_count": 84,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rich.markdown import Markdown\n",
    "Markdown(default_triage_instructions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "ca9ab99d-bc21-4cf7-a58a-261e82920566",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\n30 minute meetings are preferred, but 15 minute meetings are also acceptable.\\n'"
      ]
     },
     "execution_count": 85,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Markdown(default_cal_preferences)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "f6cd98f1-15a7-4fbb-8cce-cbbb0503d22b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"\\nUse professional and concise language. If the e-mail mentions a deadline, make sure to explicitly acknowledge and reference the deadline in your response.\\n\\nWhen responding to technical questions that require investigation:\\n- Clearly state whether you will investigate or who you will ask\\n- Provide an estimated timeline for when you'll have more information or complete the task\\n\\nWhen responding to event or conference invitations:\\n- Always acknowledge any mentioned deadlines (particularly registration deadlines)\\n- If workshops or specific topics are mentioned, ask for more specific details about them\\n- If discounts (group or early bird) are mentioned, explicitly request information about them\\n- Don't commit \\n\\nWhen responding to collaboration or project-related requests:\\n- Acknowledge any existing work or materials mentioned (drafts, slides, documents, etc.)\\n- Explicitly mention reviewing these materials before or during the meeting\\n- When scheduling meetings, clearly state the specific day, date, and time proposed\\n\\nWhen responding to meeting scheduling requests:\\n- If times are proposed, verify calendar availability for all time slots mentioned in the original email and then commit to one of the proposed times based on your availability by scheduling the meeting. Or, say you can't make it at the time proposed.\\n- If no times are proposed, then check your calendar for availability and propose multiple time options when available instead of selecting just one.\\n- Mention the meeting duration in your response to confirm you've noted it correctly.\\n- Reference the meeting's purpose in your response.\\n\""
      ]
     },
     "execution_count": 86,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Markdown(default_response_preferences)        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "d195aa00",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_memory(store, namespace, default_content=None):\n",
    "    \"\"\"Get memory from the store or initialize with default if it doesn't exist.\n",
    "    \n",
    "    Args:\n",
    "        store: LangGraph BaseStore instance to search for existing memory\n",
    "        namespace: Tuple defining the memory namespace, e.g. (\"email_assistant\", \"triage_preferences\")\n",
    "        default_content: Default content to use if memory doesn't exist\n",
    "        \n",
    "    Returns:\n",
    "        str: The content of the memory profile, either from existing memory or the default\n",
    "    \"\"\"\n",
    "    # Search for existing memory with namespace and key\n",
    "    user_preferences = store.get(namespace, \"user_preferences\")\n",
    "    \n",
    "    # If memory exists, return its content (the value)\n",
    "    if user_preferences:\n",
    "        return user_preferences.value\n",
    "    \n",
    "    # If memory doesn't exist, add it to the store and return the default content\n",
    "    else:\n",
    "        # Namespace, key, value\n",
    "        store.put(namespace, \"user_preferences\", default_content)\n",
    "        user_preferences = default_content\n",
    "    \n",
    "    # Return the default content\n",
    "    return user_preferences "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc5181e6",
   "metadata": {},
   "source": [
    "For 2) updating memory, we can use a few tricks from the [GPT-4.1 prompting guide]((https://cookbook.openai.com/examples/gpt4-1_prompting_guide)) to help us update the memory: \n",
    "\n",
    "* For optimal performance, repeat the key instructions at the start and end of the prompt\n",
    "* Create clear, explicit instructions \n",
    "* Use XML delimiters for structure\n",
    "* Provide examples "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8edb8bd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃                                               <span style=\"font-weight: bold\">Role and Objective</span>                                                ┃\n",
       "┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛\n",
       "\n",
       "You are a memory profile manager for an email assistant agent that selectively updates user preferences based on   \n",
       "feedback messages from human-in-the-loop interactions with the email assistant.                                    \n",
       "\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃                                                  <span style=\"font-weight: bold\">Instructions</span>                                                   ┃\n",
       "┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛\n",
       "\n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>NEVER overwrite the entire memory profile                                                                       \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>ONLY make targeted additions of new information                                                                 \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>ONLY update specific facts that are directly contradicted by feedback messages                                  \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>PRESERVE all other existing information in the profile                                                          \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>Format the profile consistently with the original style                                                         \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>Generate the profile as a string                                                                                \n",
       "\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃                                                 <span style=\"font-weight: bold\">Reasoning Steps</span>                                                 ┃\n",
       "┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛\n",
       "\n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 1 </span>Analyze the current memory profile structure and content                                                        \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 2 </span>Review feedback messages from human-in-the-loop interactions                                                    \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 3 </span>Extract relevant user preferences from these feedback messages (such as edits to emails/calendar invites,       \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">   </span>explicit feedback on assistant performance, user decisions to ignore certain emails)                            \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 4 </span>Compare new information against existing profile                                                                \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 5 </span>Identify only specific facts to add or update                                                                   \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 6 </span>Preserve all other existing information                                                                         \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> 7 </span>Output the complete updated profile                                                                             \n",
       "\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃                                                     <span style=\"font-weight: bold\">Example</span>                                                     ┃\n",
       "┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛\n",
       "\n",
       "&lt;memory_profile&gt; RESPOND:                                                                                          \n",
       "\n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>wife                                                                                                            \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>specific questions                                                                                              \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>system admin notifications NOTIFY:                                                                              \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>meeting invites IGNORE:                                                                                         \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>marketing emails                                                                                                \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>company-wide announcements                                                                                      \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>messages meant for other teams &lt;/memory_profile&gt;                                                                \n",
       "\n",
       "&lt;user_messages&gt; \"The assistant shouldn't have responded to that system admin notification.\" &lt;/user_messages&gt;       \n",
       "\n",
       "&lt;updated_profile&gt; RESPOND:                                                                                         \n",
       "\n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>wife                                                                                                            \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>specific questions NOTIFY:                                                                                      \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>meeting invites                                                                                                 \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>system admin notifications IGNORE:                                                                              \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>marketing emails                                                                                                \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>company-wide announcements                                                                                      \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>messages meant for other teams &lt;/updated_profile&gt;                                                               \n",
       "\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃                                     <span style=\"font-weight: bold\">Process current profile for {namespace}</span>                                     ┃\n",
       "┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛\n",
       "\n",
       "&lt;memory_profile&gt; {current_profile} &lt;/memory_profile&gt;                                                               \n",
       "\n",
       "Think step by step about what specific feedback is being provided and what specific information should be added or \n",
       "updated in the profile while preserving everything else.                                                           \n",
       "\n",
       "Think carefully and update the memory profile based upon these user messages:                                      \n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃                                               \u001b[1mRole and Objective\u001b[0m                                                ┃\n",
       "┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛\n",
       "\n",
       "You are a memory profile manager for an email assistant agent that selectively updates user preferences based on   \n",
       "feedback messages from human-in-the-loop interactions with the email assistant.                                    \n",
       "\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃                                                  \u001b[1mInstructions\u001b[0m                                                   ┃\n",
       "┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛\n",
       "\n",
       "\u001b[1;33m • \u001b[0mNEVER overwrite the entire memory profile                                                                       \n",
       "\u001b[1;33m • \u001b[0mONLY make targeted additions of new information                                                                 \n",
       "\u001b[1;33m • \u001b[0mONLY update specific facts that are directly contradicted by feedback messages                                  \n",
       "\u001b[1;33m • \u001b[0mPRESERVE all other existing information in the profile                                                          \n",
       "\u001b[1;33m • \u001b[0mFormat the profile consistently with the original style                                                         \n",
       "\u001b[1;33m • \u001b[0mGenerate the profile as a string                                                                                \n",
       "\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃                                                 \u001b[1mReasoning Steps\u001b[0m                                                 ┃\n",
       "┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛\n",
       "\n",
       "\u001b[1;33m 1 \u001b[0mAnalyze the current memory profile structure and content                                                        \n",
       "\u001b[1;33m 2 \u001b[0mReview feedback messages from human-in-the-loop interactions                                                    \n",
       "\u001b[1;33m 3 \u001b[0mExtract relevant user preferences from these feedback messages (such as edits to emails/calendar invites,       \n",
       "\u001b[1;33m   \u001b[0mexplicit feedback on assistant performance, user decisions to ignore certain emails)                            \n",
       "\u001b[1;33m 4 \u001b[0mCompare new information against existing profile                                                                \n",
       "\u001b[1;33m 5 \u001b[0mIdentify only specific facts to add or update                                                                   \n",
       "\u001b[1;33m 6 \u001b[0mPreserve all other existing information                                                                         \n",
       "\u001b[1;33m 7 \u001b[0mOutput the complete updated profile                                                                             \n",
       "\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃                                                     \u001b[1mExample\u001b[0m                                                     ┃\n",
       "┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛\n",
       "\n",
       "<memory_profile> RESPOND:                                                                                          \n",
       "\n",
       "\u001b[1;33m • \u001b[0mwife                                                                                                            \n",
       "\u001b[1;33m • \u001b[0mspecific questions                                                                                              \n",
       "\u001b[1;33m • \u001b[0msystem admin notifications NOTIFY:                                                                              \n",
       "\u001b[1;33m • \u001b[0mmeeting invites IGNORE:                                                                                         \n",
       "\u001b[1;33m • \u001b[0mmarketing emails                                                                                                \n",
       "\u001b[1;33m • \u001b[0mcompany-wide announcements                                                                                      \n",
       "\u001b[1;33m • \u001b[0mmessages meant for other teams </memory_profile>                                                                \n",
       "\n",
       "<user_messages> \"The assistant shouldn't have responded to that system admin notification.\" </user_messages>       \n",
       "\n",
       "<updated_profile> RESPOND:                                                                                         \n",
       "\n",
       "\u001b[1;33m • \u001b[0mwife                                                                                                            \n",
       "\u001b[1;33m • \u001b[0mspecific questions NOTIFY:                                                                                      \n",
       "\u001b[1;33m • \u001b[0mmeeting invites                                                                                                 \n",
       "\u001b[1;33m • \u001b[0msystem admin notifications IGNORE:                                                                              \n",
       "\u001b[1;33m • \u001b[0mmarketing emails                                                                                                \n",
       "\u001b[1;33m • \u001b[0mcompany-wide announcements                                                                                      \n",
       "\u001b[1;33m • \u001b[0mmessages meant for other teams </updated_profile>                                                               \n",
       "\n",
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃                                     \u001b[1mProcess current profile for {namespace}\u001b[0m                                     ┃\n",
       "┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛\n",
       "\n",
       "<memory_profile> {current_profile} </memory_profile>                                                               \n",
       "\n",
       "Think step by step about what specific feedback is being provided and what specific information should be added or \n",
       "updated in the profile while preserving everything else.                                                           \n",
       "\n",
       "Think carefully and update the memory profile based upon these user messages:                                      \n"
      ]
     },
     "execution_count": 90,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "Markdown(MEMORY_UPDATE_INSTRUCTIONS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "5710366b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Remember:                                                                                                          \n",
       "\n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>NEVER overwrite the entire memory profile                                                                       \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>ONLY make targeted additions of new information                                                                 \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>ONLY update specific facts that are directly contradicted by feedback messages                                  \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>PRESERVE all other existing information in the profile                                                          \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>Format the profile consistently with the original style                                                         \n",
       "<span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\"> • </span>Generate the profile as a string                                                                                \n",
       "</pre>\n"
      ],
      "text/plain": [
       "Remember:                                                                                                          \n",
       "\n",
       "\u001b[1;33m • \u001b[0mNEVER overwrite the entire memory profile                                                                       \n",
       "\u001b[1;33m • \u001b[0mONLY make targeted additions of new information                                                                 \n",
       "\u001b[1;33m • \u001b[0mONLY update specific facts that are directly contradicted by feedback messages                                  \n",
       "\u001b[1;33m • \u001b[0mPRESERVE all other existing information in the profile                                                          \n",
       "\u001b[1;33m • \u001b[0mFormat the profile consistently with the original style                                                         \n",
       "\u001b[1;33m • \u001b[0mGenerate the profile as a string                                                                                \n"
      ]
     },
     "execution_count": 91,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Markdown(MEMORY_UPDATE_INSTRUCTIONS_REINFORCEMENT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "1f8aa70e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class UserPreferences(BaseModel):\n",
    "    \"\"\"Updated user preferences based on user's feedback.\"\"\"\n",
    "    chain_of_thought: str = Field(description=\"Reasoning about which user preferences need to add / update if required\")\n",
    "    user_preferences: str = Field(description=\"Updated user preferences\")\n",
    "\n",
    "def update_memory(store, namespace, messages):\n",
    "    \"\"\"Update memory profile in the store.\n",
    "    \n",
    "    Args:\n",
    "        store: LangGraph BaseStore instance to update memory\n",
    "        namespace: Tuple defining the memory namespace, e.g. (\"email_assistant\", \"triage_preferences\")\n",
    "        messages: List of messages to update the memory with\n",
    "    \"\"\"\n",
    "\n",
    "    # Get the existing memory\n",
    "    user_preferences = store.get(namespace, \"user_preferences\")\n",
    "\n",
    "    # Update the memory\n",
    "    llm = init_chat_model(\"openai:gpt-4.1\", temperature=0.0).with_structured_output(UserPreferences)\n",
    "    result = llm.invoke(\n",
    "        [\n",
    "            {\"role\": \"system\", \"content\": MEMORY_UPDATE_INSTRUCTIONS.format(current_profile=user_preferences.value, namespace=namespace)},\n",
    "        ] + messages\n",
    "    )\n",
    "    \n",
    "    # Save the updated memory to the store\n",
    "    store.put(namespace, \"user_preferences\", result.user_preferences)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8af20960",
   "metadata": {},
   "source": [
    "We set up the triage router as we had before, with one small change"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "b1a789ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "def triage_router(state: State, store: BaseStore) -> Command[Literal[\"triage_interrupt_handler\", \"response_agent\", \"__end__\"]]:\n",
    "    \"\"\"Analyze email content to decide if we should respond, notify, or ignore.\n",
    "\n",
    "    The triage step prevents the assistant from wasting time on:\n",
    "    - Marketing emails and spam\n",
    "    - Company-wide announcements\n",
    "    - Messages meant for other teams\n",
    "    \"\"\"\n",
    "    # Parse the email input\n",
    "    author, to, subject, email_thread = parse_email(state[\"email_input\"])\n",
    "    user_prompt = triage_user_prompt.format(\n",
    "        author=author, to=to, subject=subject, email_thread=email_thread\n",
    "    )\n",
    "\n",
    "    # Create email markdown for Agent Inbox in case of notification  \n",
    "    email_markdown = format_email_markdown(subject, author, to, email_thread)\n",
    "\n",
    "    # Search for existing triage_preferences memory\n",
    "    triage_instructions = get_memory(store, (\"email_assistant\", \"triage_preferences\"), default_triage_instructions)\n",
    "\n",
    "    # Format system prompt with background and triage instructions\n",
    "    system_prompt = triage_system_prompt.format(\n",
    "        background=default_background,\n",
    "        triage_instructions=triage_instructions,\n",
    "    )\n",
    "\n",
    "    # Run the router LLM\n",
    "    result = llm_router.invoke(\n",
    "        [\n",
    "            {\"role\": \"system\", \"content\": system_prompt},\n",
    "            {\"role\": \"user\", \"content\": user_prompt},\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    # Decision\n",
    "    classification = result.classification\n",
    "\n",
    "    # Process the classification decision\n",
    "    if classification == \"respond\":\n",
    "        print(\"📧 Classification: RESPOND - This email requires a response\")\n",
    "        # Next node\n",
    "        goto = \"response_agent\"\n",
    "        # Update the state\n",
    "        update = {\n",
    "            \"classification_decision\": result.classification,\n",
    "            \"messages\": [{\"role\": \"user\",\n",
    "                            \"content\": f\"Respond to the email: {email_markdown}\"\n",
    "                        }],\n",
    "        }\n",
    "        \n",
    "    elif classification == \"ignore\":\n",
    "        print(\"🚫 Classification: IGNORE - This email can be safely ignored\")\n",
    "\n",
    "        # Next node\n",
    "        goto = END\n",
    "        # Update the state\n",
    "        update = {\n",
    "            \"classification_decision\": classification,\n",
    "        }\n",
    "\n",
    "    elif classification == \"notify\":\n",
    "        print(\"🔔 Classification: NOTIFY - This email contains important information\") \n",
    "\n",
    "        # Next node\n",
    "        goto = \"triage_interrupt_handler\"\n",
    "        # Update the state\n",
    "        update = {\n",
    "            \"classification_decision\": classification,\n",
    "        }\n",
    "\n",
    "    else:\n",
    "        raise ValueError(f\"Invalid classification: {classification}\")\n",
    "    \n",
    "    return Command(goto=goto, update=update)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6be4d63",
   "metadata": {},
   "source": [
    "We only need to make a small change to the interrupt handler to update the memory when the user provides feedback. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "f76ef46d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def triage_interrupt_handler(state: State, store: BaseStore) -> Command[Literal[\"response_agent\", \"__end__\"]]:\n",
    "    \"\"\"Handles interrupts from the triage step\"\"\"\n",
    "    \n",
    "    # Parse the email input\n",
    "    author, to, subject, email_thread = parse_email(state[\"email_input\"])\n",
    "\n",
    "    # Create email markdown for Agent Inbox in case of notification  \n",
    "    email_markdown = format_email_markdown(subject, author, to, email_thread)\n",
    "\n",
    "    # Create messages\n",
    "    messages = [{\"role\": \"user\",\n",
    "                \"content\": f\"Email to notify user about: {email_markdown}\"\n",
    "                }]\n",
    "\n",
    "    # Create interrupt for Agent Inbox\n",
    "    request = {\n",
    "        \"action_request\": {\n",
    "            \"action\": f\"Email Assistant: {state['classification_decision']}\",\n",
    "            \"args\": {}\n",
    "        },\n",
    "        \"config\": {\n",
    "            \"allow_ignore\": True,  \n",
    "            \"allow_respond\": True,\n",
    "            \"allow_edit\": False, \n",
    "            \"allow_accept\": False,  \n",
    "        },\n",
    "        # Email to show in Agent Inbox\n",
    "        \"description\": email_markdown,\n",
    "    }\n",
    "\n",
    "    # Send to Agent Inbox and wait for response\n",
    "    response = interrupt([request])[0]\n",
    "\n",
    "    # If user provides feedback, go to response agent and use feedback to respond to email   \n",
    "    if response[\"type\"] == \"response\":\n",
    "        # Add feedback to messages \n",
    "        user_input = response[\"args\"]\n",
    "        messages.append({\"role\": \"user\",\n",
    "                        \"content\": f\"User wants to reply to the email. Use this feedback to respond: {user_input}\"\n",
    "                        })\n",
    "        # This is new: update triage_preferences with feedback\n",
    "        update_memory(store, (\"email_assistant\", \"triage_preferences\"), [{\n",
    "            \"role\": \"user\",\n",
    "            \"content\": f\"The user decided to respond to the email, so update the triage preferences to capture this.\"\n",
    "        }] + messages)\n",
    "\n",
    "        goto = \"response_agent\"\n",
    "\n",
    "    # If user ignores email, go to END\n",
    "    elif response[\"type\"] == \"ignore\":\n",
    "        # Make note of the user's decision to ignore the email\n",
    "        messages.append({\"role\": \"user\",\n",
    "                        \"content\": f\"The user decided to ignore the email even though it was classified as notify. Update triage preferences to capture this.\"\n",
    "                        })\n",
    "        # This is new: triage_preferences with feedback\n",
    "        update_memory(store, (\"email_assistant\", \"triage_preferences\"), messages)\n",
    "        goto = END\n",
    "\n",
    "    # Catch all other responses\n",
    "    else:\n",
    "        raise ValueError(f\"Invalid response: {response}\")\n",
    "\n",
    "    # Update the state \n",
    "    update = {\n",
    "        \"messages\": messages,\n",
    "    }\n",
    "\n",
    "    return Command(goto=goto, update=update)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cd428f5",
   "metadata": {},
   "source": [
    "### Incorporating Memory into LLM Responses\n",
    "\n",
    "Now that we have memory managers set up, we can use the stored preferences when generating responses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "a82b17a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def llm_call(state: State, store: BaseStore):\n",
    "    \"\"\"LLM decides whether to call a tool or not\"\"\"\n",
    "\n",
    "    # Search for existing cal_preferences memory\n",
    "    cal_preferences = get_memory(store, (\"email_assistant\", \"cal_preferences\"), default_cal_preferences)\n",
    "    \n",
    "    # Search for existing response_preferences memory\n",
    "    response_preferences = get_memory(store, (\"email_assistant\", \"response_preferences\"), default_response_preferences)\n",
    "\n",
    "    return {\n",
    "        \"messages\": [\n",
    "            llm_with_tools.invoke(\n",
    "                [\n",
    "                    {\"role\": \"system\", \"content\": agent_system_prompt_hitl_memory.format(tools_prompt=HITL_MEMORY_TOOLS_PROMPT,\n",
    "                                                                                         background=default_background,\n",
    "                                                                                         response_preferences=response_preferences, \n",
    "                                                                                         cal_preferences=cal_preferences)}\n",
    "                ]\n",
    "                + state[\"messages\"]\n",
    "            )\n",
    "        ]\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e60aff5d",
   "metadata": {},
   "source": [
    "### Memory Integration in the Interrupt Handler\n",
    "\n",
    "Similarly, we'll add memory to the interrupt handler! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "126d3680",
   "metadata": {},
   "outputs": [],
   "source": [
    "def interrupt_handler(state: State, store: BaseStore) -> Command[Literal[\"llm_call\", \"__end__\"]]:\n",
    "    \"\"\"Creates an interrupt for human review of tool calls\"\"\"\n",
    "    \n",
    "    # Store messages\n",
    "    result = []\n",
    "\n",
    "    # Go to the LLM call node next\n",
    "    goto = \"llm_call\"\n",
    "\n",
    "    # Iterate over the tool calls in the last message\n",
    "    for tool_call in state[\"messages\"][-1].tool_calls:\n",
    "        \n",
    "        # Allowed tools for HITL\n",
    "        hitl_tools = [\"write_email\", \"schedule_meeting\", \"Question\"]\n",
    "        \n",
    "        # If tool is not in our HITL list, execute it directly without interruption\n",
    "        if tool_call[\"name\"] not in hitl_tools:\n",
    "\n",
    "            # Execute tool without interruption\n",
    "            tool = tools_by_name[tool_call[\"name\"]]\n",
    "            observation = tool.invoke(tool_call[\"args\"])\n",
    "            result.append({\"role\": \"tool\", \"content\": observation, \"tool_call_id\": tool_call[\"id\"]})\n",
    "            continue\n",
    "            \n",
    "        # Get original email from email_input in state\n",
    "        email_input = state[\"email_input\"]\n",
    "        author, to, subject, email_thread = parse_email(email_input)\n",
    "        original_email_markdown = format_email_markdown(subject, author, to, email_thread)\n",
    "        \n",
    "        # Format tool call for display and prepend the original email\n",
    "        tool_display = format_for_display(tool_call)\n",
    "        description = original_email_markdown + tool_display\n",
    "\n",
    "        # Configure what actions are allowed in Agent Inbox\n",
    "        if tool_call[\"name\"] == \"write_email\":\n",
    "            config = {\n",
    "                \"allow_ignore\": True,\n",
    "                \"allow_respond\": True,\n",
    "                \"allow_edit\": True,\n",
    "                \"allow_accept\": True,\n",
    "            }\n",
    "        elif tool_call[\"name\"] == \"schedule_meeting\":\n",
    "            config = {\n",
    "                \"allow_ignore\": True,\n",
    "                \"allow_respond\": True,\n",
    "                \"allow_edit\": True,\n",
    "                \"allow_accept\": True,\n",
    "            }\n",
    "        elif tool_call[\"name\"] == \"Question\":\n",
    "            config = {\n",
    "                \"allow_ignore\": True,\n",
    "                \"allow_respond\": True,\n",
    "                \"allow_edit\": False,\n",
    "                \"allow_accept\": False,\n",
    "            }\n",
    "        else:\n",
    "            raise ValueError(f\"Invalid tool call: {tool_call['name']}\")\n",
    "\n",
    "        # Create the interrupt request\n",
    "        request = {\n",
    "            \"action_request\": {\n",
    "                \"action\": tool_call[\"name\"],\n",
    "                \"args\": tool_call[\"args\"]\n",
    "            },\n",
    "            \"config\": config,\n",
    "            \"description\": description,\n",
    "        }\n",
    "\n",
    "        # Send to Agent Inbox and wait for response\n",
    "        response = interrupt([request])[0]\n",
    "\n",
    "        # Handle the responses \n",
    "        if response[\"type\"] == \"accept\":\n",
    "\n",
    "            # Execute the tool with original args\n",
    "            tool = tools_by_name[tool_call[\"name\"]]\n",
    "            observation = tool.invoke(tool_call[\"args\"])\n",
    "            result.append({\"role\": \"tool\", \"content\": observation, \"tool_call_id\": tool_call[\"id\"]})\n",
    "                        \n",
    "        elif response[\"type\"] == \"edit\":\n",
    "\n",
    "            # Tool selection \n",
    "            tool = tools_by_name[tool_call[\"name\"]]\n",
    "            initial_tool_call = tool_call[\"args\"]\n",
    "            \n",
    "            # Get edited args from Agent Inbox\n",
    "            edited_args = response[\"args\"][\"args\"]\n",
    "\n",
    "            # Update the AI message's tool call with edited content (reference to the message in the state)\n",
    "            ai_message = state[\"messages\"][-1] # Get the most recent message from the state\n",
    "            current_id = tool_call[\"id\"] # Store the ID of the tool call being edited\n",
    "            \n",
    "            # Create a new list of tool calls by filtering out the one being edited and adding the updated version\n",
    "            # This avoids modifying the original list directly (immutable approach)\n",
    "            updated_tool_calls = [tc for tc in ai_message.tool_calls if tc[\"id\"] != current_id] + [\n",
    "                {\"type\": \"tool_call\", \"name\": tool_call[\"name\"], \"args\": edited_args, \"id\": current_id}\n",
    "            ]\n",
    "\n",
    "            # Create a new copy of the message with updated tool calls rather than modifying the original\n",
    "            # This ensures state immutability and prevents side effects in other parts of the code\n",
    "            # When we update the messages state key (\"messages\": result), the add_messages reducer will\n",
    "            # overwrite existing messages by id and we take advantage of this here to update the tool calls.\n",
    "            result.append(ai_message.model_copy(update={\"tool_calls\": updated_tool_calls}))\n",
    "\n",
    "            # Save feedback in memory and update the write_email tool call with the edited content from Agent Inbox\n",
    "            if tool_call[\"name\"] == \"write_email\":\n",
    "                \n",
    "                # Execute the tool with edited args\n",
    "                observation = tool.invoke(edited_args)\n",
    "                \n",
    "                # Add only the tool response message\n",
    "                result.append({\"role\": \"tool\", \"content\": observation, \"tool_call_id\": current_id})\n",
    "\n",
    "                # This is new: update the memory\n",
    "                update_memory(store, (\"email_assistant\", \"response_preferences\"), [{\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": f\"User edited the email response. Here is the initial email generated by the assistant: {initial_tool_call}. Here is the edited email: {edited_args}. Follow all instructions above, and remember: {MEMORY_UPDATE_INSTRUCTIONS_REINFORCEMENT}.\"\n",
    "                }])\n",
    "            \n",
    "            # Save feedback in memory and update the schedule_meeting tool call with the edited content from Agent Inbox\n",
    "            elif tool_call[\"name\"] == \"schedule_meeting\":\n",
    "                \n",
    "                # Execute the tool with edited args\n",
    "                observation = tool.invoke(edited_args)\n",
    "                \n",
    "                # Add only the tool response message\n",
    "                result.append({\"role\": \"tool\", \"content\": observation, \"tool_call_id\": current_id})\n",
    "\n",
    "                # This is new: update the memory\n",
    "                update_memory(store, (\"email_assistant\", \"cal_preferences\"), [{\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": f\"User edited the calendar invitation. Here is the initial calendar invitation generated by the assistant: {initial_tool_call}. Here is the edited calendar invitation: {edited_args}. Follow all instructions above, and remember: {MEMORY_UPDATE_INSTRUCTIONS_REINFORCEMENT}.\"\n",
    "                }])\n",
    "            \n",
    "            # Catch all other tool calls\n",
    "            else:\n",
    "                raise ValueError(f\"Invalid tool call: {tool_call['name']}\")\n",
    "\n",
    "        elif response[\"type\"] == \"ignore\":\n",
    "\n",
    "            if tool_call[\"name\"] == \"write_email\":\n",
    "                # Don't execute the tool, and tell the agent how to proceed\n",
    "                result.append({\"role\": \"tool\", \"content\": \"User ignored this email draft. Ignore this email and end the workflow.\", \"tool_call_id\": tool_call[\"id\"]})\n",
    "                # Go to END\n",
    "                goto = END\n",
    "                # This is new: update the memory\n",
    "                update_memory(store, (\"email_assistant\", \"triage_preferences\"), state[\"messages\"] + result + [{\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": f\"The user ignored the email draft. That means they did not want to respond to the email. Update the triage preferences to ensure emails of this type are not classified as respond. Follow all instructions above, and remember: {MEMORY_UPDATE_INSTRUCTIONS_REINFORCEMENT}.\"\n",
    "                }])\n",
    "\n",
    "            elif tool_call[\"name\"] == \"schedule_meeting\":\n",
    "                # Don't execute the tool, and tell the agent how to proceed\n",
    "                result.append({\"role\": \"tool\", \"content\": \"User ignored this calendar meeting draft. Ignore this email and end the workflow.\", \"tool_call_id\": tool_call[\"id\"]})\n",
    "                # Go to END\n",
    "                goto = END\n",
    "                # This is new: update the memory\n",
    "                update_memory(store, (\"email_assistant\", \"triage_preferences\"), state[\"messages\"] + result + [{\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": f\"The user ignored the calendar meeting draft. That means they did not want to schedule a meeting for this email. Update the triage preferences to ensure emails of this type are not classified as respond. Follow all instructions above, and remember: {MEMORY_UPDATE_INSTRUCTIONS_REINFORCEMENT}.\"\n",
    "                }])\n",
    "\n",
    "            elif tool_call[\"name\"] == \"Question\":\n",
    "                # Don't execute the tool, and tell the agent how to proceed\n",
    "                result.append({\"role\": \"tool\", \"content\": \"User ignored this question. Ignore this email and end the workflow.\", \"tool_call_id\": tool_call[\"id\"]})\n",
    "                # Go to END\n",
    "                goto = END\n",
    "                # This is new: update the memory\n",
    "                update_memory(store, (\"email_assistant\", \"triage_preferences\"), state[\"messages\"] + result + [{\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": f\"The user ignored the Question. That means they did not want to answer the question or deal with this email. Update the triage preferences to ensure emails of this type are not classified as respond. Follow all instructions above, and remember: {MEMORY_UPDATE_INSTRUCTIONS_REINFORCEMENT}.\"\n",
    "                }])\n",
    "\n",
    "            else:\n",
    "                raise ValueError(f\"Invalid tool call: {tool_call['name']}\")\n",
    "\n",
    "        elif response[\"type\"] == \"response\":\n",
    "            # User provided feedback\n",
    "            user_feedback = response[\"args\"]\n",
    "            if tool_call[\"name\"] == \"write_email\":\n",
    "                # Don't execute the tool, and add a message with the user feedback to incorporate into the email\n",
    "                result.append({\"role\": \"tool\", \"content\": f\"User gave feedback, which can we incorporate into the email. Feedback: {user_feedback}\", \"tool_call_id\": tool_call[\"id\"]})\n",
    "                # This is new: update the memory\n",
    "                update_memory(store, (\"email_assistant\", \"response_preferences\"), state[\"messages\"] + result + [{\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": f\"User gave feedback, which we can use to update the response preferences. Follow all instructions above, and remember: {MEMORY_UPDATE_INSTRUCTIONS_REINFORCEMENT}.\"\n",
    "                }])\n",
    "\n",
    "            elif tool_call[\"name\"] == \"schedule_meeting\":\n",
    "                # Don't execute the tool, and add a message with the user feedback to incorporate into the email\n",
    "                result.append({\"role\": \"tool\", \"content\": f\"User gave feedback, which can we incorporate into the meeting request. Feedback: {user_feedback}\", \"tool_call_id\": tool_call[\"id\"]})\n",
    "                # This is new: update the memory\n",
    "                update_memory(store, (\"email_assistant\", \"cal_preferences\"), state[\"messages\"] + result + [{\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": f\"User gave feedback, which we can use to update the calendar preferences. Follow all instructions above, and remember: {MEMORY_UPDATE_INSTRUCTIONS_REINFORCEMENT}.\"\n",
    "                }])\n",
    "\n",
    "            elif tool_call[\"name\"] == \"Question\":\n",
    "                # Don't execute the tool, and add a message with the user feedback to incorporate into the email\n",
    "                result.append({\"role\": \"tool\", \"content\": f\"User answered the question, which can we can use for any follow up actions. Feedback: {user_feedback}\", \"tool_call_id\": tool_call[\"id\"]})\n",
    "\n",
    "            else:\n",
    "                raise ValueError(f\"Invalid tool call: {tool_call['name']}\")\n",
    "\n",
    "    # Update the state \n",
    "    update = {\n",
    "        \"messages\": result,\n",
    "    }\n",
    "\n",
    "    return Command(goto=goto, update=update)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecedcaec",
   "metadata": {},
   "source": [
    "The rest is the same as before!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7041f50d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from email_assistant.utils import show_graph\n",
    "\n",
    "# Conditional edge function\n",
    "def should_continue(state: State, store: BaseStore) -> Literal[\"interrupt_handler\", END]:\n",
    "    \"\"\"Route to tool handler, or end if Done tool called\"\"\"\n",
    "    messages = state[\"messages\"]\n",
    "    last_message = messages[-1]\n",
    "    if last_message.tool_calls:\n",
    "        for tool_call in last_message.tool_calls: \n",
    "            if tool_call[\"name\"] == \"Done\":\n",
    "                return END\n",
    "            else:\n",
    "                return \"interrupt_handler\"\n",
    "\n",
    "# Build workflow\n",
    "agent_builder = StateGraph(State)\n",
    "\n",
    "# Add nodes - with store parameter\n",
    "agent_builder.add_node(\"llm_call\", llm_call)\n",
    "agent_builder.add_node(\"interrupt_handler\", interrupt_handler)\n",
    "\n",
    "# Add edges\n",
    "agent_builder.add_edge(START, \"llm_call\")\n",
    "agent_builder.add_conditional_edges(\n",
    "    \"llm_call\",\n",
    "    should_continue,\n",
    "    {\n",
    "        \"interrupt_handler\": \"interrupt_handler\",\n",
    "        END: END,\n",
    "    },\n",
    ")\n",
    "\n",
    "# Compile the agent\n",
    "response_agent = agent_builder.compile()\n",
    "\n",
    "# Build overall workflow with store and checkpointer\n",
    "overall_workflow = (\n",
    "    StateGraph(State, input=StateInput)\n",
    "    .add_node(triage_router)\n",
    "    .add_node(triage_interrupt_handler)\n",
    "    .add_node(\"response_agent\", response_agent)\n",
    "    .add_edge(START, \"triage_router\")\n",
    ")\n",
    "\n",
    "email_assistant = overall_workflow.compile()\n",
    "show_graph(email_assistant)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43747219",
   "metadata": {},
   "source": [
    "## Testing the agent with memory\n",
    "\n",
    "Now that we've implemented memory into our email assistant, let's test how the system learns from user feedback and adapts over time. This testing section explores how different types of user interactions create distinct memory updates that improve the assistant's future performance.\n",
    "\n",
    "The key questions we're answering through these tests:\n",
    "1. How does the system capture and store user preferences?\n",
    "2. How do these stored preferences affect future decisions?\n",
    "3. What patterns of interaction lead to which types of memory updates?\n",
    "\n",
    "First, let's build a helper function to display memory content so we can track how it evolves throughout our tests:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "59079929",
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid \n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.types import Command\n",
    "from langgraph.store.memory import InMemoryStore\n",
    "\n",
    "# Helper function to display memory content\n",
    "def display_memory_content(store, namespace=None):\n",
    "    # Display current memory content for all namespaces\n",
    "    print(\"\\n======= CURRENT MEMORY CONTENT =======\")\n",
    "    if namespace:\n",
    "        memory = store.get(namespace, \"user_preferences\")\n",
    "        if memory:\n",
    "            print(f\"\\n--- {namespace[1]} ---\")\n",
    "            print(memory.value)\n",
    "        else:\n",
    "            print(f\"\\n--- {namespace[1]} ---\")\n",
    "            print(\"No memory found\")\n",
    "    else:\n",
    "        for namespace in [\n",
    "            (\"email_assistant\", \"triage_preferences\"),\n",
    "            (\"email_assistant\", \"response_preferences\"),\n",
    "            (\"email_assistant\", \"cal_preferences\"),\n",
    "            (\"email_assistant\", \"background\")\n",
    "        ]:\n",
    "            memory = store.get(namespace, \"user_preferences\")\n",
    "            if memory:\n",
    "                print(f\"\\n--- {namespace[1]} ---\")\n",
    "                print(memory.value)\n",
    "            else:\n",
    "                print(f\"\\n--- {namespace[1]} ---\")\n",
    "                print(\"No memory found\")\n",
    "            print(\"=======================================\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "397114bf",
   "metadata": {},
   "source": [
    "### Accept `write_email` and `schedule_meeting`\n",
    "\n",
    "Our first test examines what happens when a user accepts the agent's actions without modification. This baseline case helps us understand how the system behaves when no feedback is provided:\n",
    "\n",
    "1. We'll use the same tax planning email from our previous tests\n",
    "2. The system will classify it as \"RESPOND\" and propose scheduling a meeting\n",
    "3. We'll accept the meeting schedule without changes\n",
    "4. The agent will generate an email confirming the meeting\n",
    "5. We'll accept the email without changes\n",
    "\n",
    "This test demonstrates the default behavior of our memory-enabled system. When a user simply accepts proposed actions, we expect minimal or no memory updates since there's no explicit feedback to learn from. However, the system will still leverage existing memory (if any) when generating its responses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "649cee4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Respond - Meeting Request Email\n",
    "email_input_respond = {\n",
    "    \"to\": \"Lance Martin <lance@company.com>\",\n",
    "    \"author\": \"Project Manager <pm@client.com>\",\n",
    "    \"subject\": \"Tax season let's schedule call\",\n",
    "    \"email_thread\": \"Lance,\\n\\nIt's tax season again, and I wanted to schedule a call to discuss your tax planning strategies for this year. I have some suggestions that could potentially save you money.\\n\\nAre you available sometime next week? Tuesday or Thursday afternoon would work best for me, for about 45 minutes.\\n\\nRegards,\\nProject Manager\"\n",
    "}\n",
    "\n",
    "# Compile the graph\n",
    "checkpointer = MemorySaver()\n",
    "store = InMemoryStore()\n",
    "graph = overall_workflow.compile(checkpointer=checkpointer, store=store)\n",
    "thread_id_1 = uuid.uuid4()\n",
    "thread_config_1 = {\"configurable\": {\"thread_id\": thread_id_1}}\n",
    "\n",
    "# Run the graph until the first interrupt \n",
    "# Email will be classified as \"respond\" \n",
    "# Agent will create a schedule_meeting and write_email tool call\n",
    "print(\"Running the graph until the first interrupt...\")\n",
    "for chunk in graph.stream({\"email_input\": email_input_respond}, config=thread_config_1):\n",
    "    # Inspect interrupt object if present\n",
    "    if '__interrupt__' in chunk:\n",
    "        Interrupt_Object = chunk['__interrupt__'][0]\n",
    "        print(\"\\nINTERRUPT OBJECT:\")\n",
    "        print(f\"Action Request: {Interrupt_Object.value[0]['action_request']}\")\n",
    "\n",
    "# Check memory after first interrupt\n",
    "display_memory_content(store)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "878e199e",
   "metadata": {},
   "source": [
    "Accept the `schedule_meeting` tool call\n",
    "\n",
    "As we examine the initial `schedule_meeting` proposal, note how the system uses existing memory to inform its decisions:\n",
    "\n",
    "1. The default calendar preferences show a preference for 30-minute meetings, though the email requests 45 minutes\n",
    "2. The agent still proposes a 45-minute meeting, respecting the sender's specific request\n",
    "3. We accept this proposal without modification to see if simple acceptance triggers any memory updates\n",
    "\n",
    "After running this step, we'll check the memory contents to confirm whether acceptance alone leads to memory updates. Simple acceptance represents the baseline user experience - the system works as intended without requiring adjustments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9589423b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"\\nSimulating user accepting the {Interrupt_Object.value[0]['action_request']['action']} tool call...\")\n",
    "for chunk in graph.stream(Command(resume=[{\"type\": \"accept\"}]), config=thread_config_1):\n",
    "    # Inspect interrupt object if present\n",
    "    if '__interrupt__' in chunk:\n",
    "        Interrupt_Object = chunk['__interrupt__'][0]\n",
    "        print(\"\\nINTERRUPT OBJECT:\")\n",
    "        print(f\"Action Request: {Interrupt_Object.value[0]['action_request']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6b80f99",
   "metadata": {},
   "source": [
    "Accept the `write_email` tool call\n",
    "\n",
    "Now we'll accept the email draft that confirms the meeting scheduling:\n",
    "\n",
    "1. The email draft is generated with knowledge of our calendar preferences\n",
    "2. It includes details about the meeting time, duration, and purpose\n",
    "3. We'll accept it without changes to complete the baseline test case\n",
    "\n",
    "After accepting, we'll check all memory stores to see if any updates occurred. As expected, simply accepting the agent's proposals doesn't provide strong learning signals - there's no clear feedback about what the user likes or dislikes about the agent's approach.\n",
    "\n",
    "The trace link shows the complete workflow execution, where we can see that the memory is used in the LLM call for response generation, but no memory updates occur, which is the expected behavior for simple acceptances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12035cf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"\\nSimulating user accepting the {Interrupt_Object.value[0]['action_request']['action']} tool call...\")\n",
    "for chunk in graph.stream(Command(resume=[{\"type\": \"accept\"}]), config=thread_config_1):\n",
    "    # Inspect response_agent most recent message\n",
    "    if 'response_agent' in chunk:\n",
    "        chunk['response_agent']['messages'][-1].pretty_print()\n",
    "    # Inspect interrupt object if present\n",
    "    if '__interrupt__' in chunk:\n",
    "        Interrupt_Object = chunk['__interrupt__'][0]\n",
    "        print(\"\\nINTERRUPT OBJECT:\")\n",
    "        print(f\"Action Request: {Interrupt_Object.value[0]['action_request']}\")\n",
    "\n",
    "# Check memory after accepting the write_email tool call\n",
    "display_memory_content(store)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcbc178d",
   "metadata": {},
   "source": [
    "We can look at the full messages, and the trace: \n",
    "\n",
    "https://smith.langchain.com/public/86ff6474-29fe-452e-8829-b05a91b458eb/r\n",
    "\n",
    "You'll notice that memory is used in the LLM call to respond. \n",
    "\n",
    "But the memory store is *not* updated, because we haven't added any feedback via HITL."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10ce8197",
   "metadata": {},
   "outputs": [],
   "source": [
    "state = graph.get_state(thread_config_1)\n",
    "for m in state.values['messages']:\n",
    "    m.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58201a21",
   "metadata": {},
   "source": [
    "### Edit `write_email` and `schedule_meeting`\n",
    "\n",
    "This test explores how the system learns from direct edits to its proposed actions. When users modify the agent's suggestions, it creates clear, specific learning signals about their preferences:\n",
    "\n",
    "1. We'll use the same tax planning email as before\n",
    "2. When the agent proposes a 45-minute meeting, we'll edit it to:\n",
    "   - Change the duration to 30 minutes (matching our stored preference)\n",
    "   - Make the subject line more concise\n",
    "3. When the agent drafts an email, we'll edit it to be:\n",
    "   - Shorter and less formal\n",
    "   - Structured differently\n",
    "\n",
    "Edits provide the most explicit feedback about user preferences, letting the system learn exactly what changes are desired. We expect to see specific, targeted updates to our memory stores that reflect these edits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac260423",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Same email as before\n",
    "email_input_respond = {\n",
    "    \"to\": \"Lance Martin <lance@company.com>\",\n",
    "    \"author\": \"Project Manager <pm@client.com>\",\n",
    "    \"subject\": \"Tax season let's schedule call\",\n",
    "    \"email_thread\": \"Lance,\\n\\nIt's tax season again, and I wanted to schedule a call to discuss your tax planning strategies for this year. I have some suggestions that could potentially save you money.\\n\\nAre you available sometime next week? Tuesday or Thursday afternoon would work best for me, for about 45 minutes.\\n\\nRegards,\\nProject Manager\"\n",
    "}\n",
    "\n",
    "# Compile the graph with new thread\n",
    "checkpointer = MemorySaver()\n",
    "store = InMemoryStore()\n",
    "graph = overall_workflow.compile(checkpointer=checkpointer, store=store)\n",
    "thread_id_2 = uuid.uuid4()\n",
    "thread_config_2 = {\"configurable\": {\"thread_id\": thread_id_2}}\n",
    "\n",
    "# Run the graph until the first interrupt - will be classified as \"respond\" and the agent will create a write_email tool call\n",
    "print(\"Running the graph until the first interrupt...\")\n",
    "for chunk in graph.stream({\"email_input\": email_input_respond}, config=thread_config_2):\n",
    "    # Inspect interrupt object if present\n",
    "    if '__interrupt__' in chunk:\n",
    "        Interrupt_Object = chunk['__interrupt__'][0]\n",
    "        print(\"\\nINTERRUPT OBJECT:\")\n",
    "        print(f\"Action Request: {Interrupt_Object.value[0]['action_request']}\")\n",
    "\n",
    "# Check memory after first interrupt\n",
    "display_memory_content(store,(\"email_assistant\", \"cal_preferences\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d73ba71",
   "metadata": {},
   "source": [
    "Edit the `schedule_meeting` tool call\n",
    "\n",
    "When we edit the meeting proposal, we're providing direct, explicit feedback about our preferences. This creates a significant learning opportunity for the system:\n",
    "\n",
    "1. The agent initially proposes a 45-minute meeting (the duration requested in the email)\n",
    "2. We edit it to 30 minutes and simplify the subject from \"Tax Planning Strategies Discussion\" to \"Tax Planning Discussion\"\n",
    "3. This creates clear, specific feedback about our time preferences and naming conventions\n",
    "\n",
    "After the edit, we'll check the calendar preferences memory store to see how it's updated. The memory update should capture both:\n",
    "- Our preference for shorter 30-minute meetings\n",
    "- Our preference for more concise meeting subjects\n",
    "\n",
    "The trace reveals the precise memory update logic, showing how the system analyzes the difference between its proposal and our edits to extract meaningful patterns and preferences. We can see the detailed justification for each memory update, ensuring transparency in the learning process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af760977",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now simulate user editing the schedule_meeting tool call\n",
    "print(\"\\nSimulating user editing the schedule_meeting tool call...\")\n",
    "edited_schedule_args = {\n",
    "    \"attendees\": [\"pm@client.com\", \"lance@company.com\"],\n",
    "    \"subject\": \"Tax Planning Discussion\",\n",
    "    \"duration_minutes\": 30, # Changed from 45 to 30\n",
    "    \"preferred_day\": \"2025-04-22\",\n",
    "    \"start_time\": 14 \n",
    "}\n",
    "for chunk in graph.stream(Command(resume=[{\"type\": \"edit\", \"args\": {\"args\": edited_schedule_args}}]), config=thread_config_2):\n",
    "    # Inspect response_agent most recent message\n",
    "    if 'response_agent' in chunk:\n",
    "        chunk['response_agent']['messages'][-1].pretty_print()\n",
    "    # Inspect interrupt object if present\n",
    "    if '__interrupt__' in chunk:\n",
    "        Interrupt_Object = chunk['__interrupt__'][0]\n",
    "        print(\"\\nINTERRUPT OBJECT:\")\n",
    "        print(f\"Action Request: {Interrupt_Object.value[0]['action_request']}\")\n",
    "\n",
    "# Check memory after editing schedule_meeting\n",
    "print(\"\\nChecking memory after editing schedule_meeting:\")\n",
    "display_memory_content(store,(\"email_assistant\", \"cal_preferences\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbbb324f",
   "metadata": {
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "```\n",
    "{'preferences': '\\n30 minute meetings are preferred, but 15 minute meetings are also acceptable.\\n'}\n",
    "```\n",
    "\n",
    "```\n",
    "{'preferences': \"30 minute meetings are preferred, but 15 minute meetings are also acceptable.\\n\\nUser prefers 30 minute meetings over longer durations such as 45 minutes. When scheduling, default to 30 minutes unless otherwise specified. Subject lines should be concise (e.g., 'Tax Planning Discussion' instead of 'Tax Planning Strategies Discussion').\"}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0dfc585a",
   "metadata": {},
   "source": [
    "Looking at the memory after editing the calendar invitation, we can see that it's been updated:\n",
    "\n",
    "1. The system has identified that we prefer 30-minute meetings over longer durations\n",
    "2. It's also captured our preference for concise meeting subjects\n",
    "\n",
    "What's particularly impressive about this memory update is:\n",
    "- It doesn't just record our specific edit, but generalizes to a broader preference pattern\n",
    "- It preserves all existing memory content while adding the new information\n",
    "- It extracts multiple preference signals from a single edit interaction\n",
    "\n",
    "Now, let's edit the email draft to see how the system captures different types of communication preferences:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81a1fa37",
   "metadata": {},
   "outputs": [],
   "source": [
    "display_memory_content(store,(\"email_assistant\", \"response_preferences\"))\n",
    "# Now simulate user editing the write_email tool call\n",
    "print(\"\\nSimulating user editing the write_email tool call...\")\n",
    "edited_email_args = {\n",
    "    \"to\": \"pm@client.com\",\n",
    "    \"subject\": \"Re: Tax season let's schedule call\",\n",
    "    \"content\": \"Thanks! I scheduled a 30-minute call next Thursday at 3:00 PM. Would that work for you?\\n\\nBest regards,\\nLance Martin\"\n",
    "}\n",
    "for chunk in graph.stream(Command(resume=[{\"type\": \"edit\", \"args\": {\"args\": edited_email_args}}]), config=thread_config_2):\n",
    "    # Inspect response_agent most recent message\n",
    "    if 'response_agent' in chunk:\n",
    "        chunk['response_agent']['messages'][-1].pretty_print()\n",
    "    # Inspect interrupt object if present\n",
    "    if '__interrupt__' in chunk:\n",
    "        Interrupt_Object = chunk['__interrupt__'][0]\n",
    "        print(\"\\nINTERRUPT OBJECT:\")\n",
    "        print(f\"Action Request: {Interrupt_Object.value[0]['action_request']}\")\n",
    "\n",
    "# Check memory after editing write_email\n",
    "print(\"\\nChecking memory after editing write_email:\")\n",
    "display_memory_content(store,(\"email_assistant\", \"response_preferences\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ffbd5f9",
   "metadata": {},
   "source": [
    "Our email edit reveals even more sophisticated learning capabilities:\n",
    "\n",
    "1. We've dramatically shortened and simplified the email content\n",
    "2. We've changed the tone to be more casual\n",
    "3. We've added a question asking for confirmation rather than assuming the time works\n",
    "4. We've slightly altered the meeting details (day and time)\n",
    "\n",
    "Looking at the updated memory, we can see that the system has extracted a key insight about our communication style:\n",
    "\n",
    "```\n",
    "When scheduling a meeting, ask the recipient to confirm if the proposed time works for them, rather than assuming and stating the meeting is already scheduled.\n",
    "```\n",
    "\n",
    "This demonstrates the system's ability to:\n",
    "- Analyze our edit not just at a superficial level, but to understand intent\n",
    "- Extract generalizable principles from specific examples\n",
    "- Preserve all existing guidance while adding new insights\n",
    "- Maintain the organization and structure of the memory\n",
    "\n",
    "These targeted, high-quality memory updates will improve all future interactions without requiring repeated corrections."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ad818d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "state = graph.get_state(thread_config_2)\n",
    "for m in state.values['messages']:\n",
    "    m.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d92a42b",
   "metadata": {},
   "source": [
    "### Respond (with feedback) `write_email` and `schedule_meeting`\n",
    "\n",
    "Our final test set explores the \"response\" feedback pattern - providing guidance without directly editing or accepting. This conversational feedback mechanism offers a middle ground between acceptance and editing:\n",
    "\n",
    "1. First, we'll test feedback for meeting scheduling by requesting:\n",
    "   - Shorter duration (30 minutes instead of 45)\n",
    "   - Afternoon meeting times (after 2pm)\n",
    "   \n",
    "2. Next, we'll test feedback for email drafting by requesting:\n",
    "   - Shorter, less formal language\n",
    "   - A specific closing statement about looking forward to the meeting\n",
    "   \n",
    "3. Finally, we'll test feedback for questions by providing:\n",
    "   - A direct answer with additional context\n",
    "   - Specific preferences (brunch location, time)\n",
    "\n",
    "This natural language feedback approach lets users guide the assistant without having to do the work themselves. We expect to see detailed memory updates that extract the general principles from our specific feedback."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07676231",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Respond - Meeting Request Email\n",
    "email_input_respond = {\n",
    "    \"to\": \"Lance Martin <lance@company.com>\",\n",
    "    \"author\": \"Project Manager <pm@client.com>\",\n",
    "    \"subject\": \"Tax season let's schedule call\",\n",
    "    \"email_thread\": \"Lance,\\n\\nIt's tax season again, and I wanted to schedule a call to discuss your tax planning strategies for this year. I have some suggestions that could potentially save you money.\\n\\nAre you available sometime next week? Tuesday or Thursday afternoon would work best for me, for about 45 minutes.\\n\\nRegards,\\nProject Manager\"\n",
    "}\n",
    "\n",
    "# Compile the graph\n",
    "checkpointer = MemorySaver()\n",
    "store = InMemoryStore()\n",
    "graph = overall_workflow.compile(checkpointer=checkpointer, store=store)\n",
    "thread_id_5 = uuid.uuid4()\n",
    "thread_config_5 = {\"configurable\": {\"thread_id\": thread_id_5}}\n",
    "\n",
    "# Run the graph until the first interrupt \n",
    "# Email will be classified as \"respond\" \n",
    "# Agent will create a schedule_meeting and write_email tool call\n",
    "print(\"Running the graph until the first interrupt...\")\n",
    "for chunk in graph.stream({\"email_input\": email_input_respond}, config=thread_config_5):\n",
    "    # Inspect interrupt object if present\n",
    "    if '__interrupt__' in chunk:\n",
    "        Interrupt_Object = chunk['__interrupt__'][0]\n",
    "        print(\"\\nINTERRUPT OBJECT:\")\n",
    "        print(f\"Action Request: {Interrupt_Object.value[0]['action_request']}\")\n",
    "\n",
    "# Check memory after first interrupt \n",
    "display_memory_content(store, (\"email_assistant\", \"cal_preferences\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b85fc45d",
   "metadata": {},
   "source": [
    "Provide feedback for the `schedule_meeting` tool call\n",
    "\n",
    "Instead of directly editing the meeting proposal or simply accepting it, we'll provide natural language feedback:\n",
    "\n",
    "1. We request a 30-minute meeting instead of 45 minutes\n",
    "2. We express a preference for afternoon meetings after 2pm\n",
    "3. The system must interpret this feedback and generate a new proposal\n",
    "\n",
    "This conversational approach is often more natural and efficient than direct editing, especially for mobile users or those who prefer to give high-level direction rather than detailed edits.\n",
    "\n",
    "After providing feedback, we'll examine the calendar preferences memory to see how this natural language guidance is captured. We expect to see the system extract both the meeting duration and time-of-day preferences as general principles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30a151f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"\\nSimulating user providing feedback for the {Interrupt_Object.value[0]['action_request']['action']} tool call...\")\n",
    "for chunk in graph.stream(Command(resume=[{\"type\": \"response\", \"args\": \"Please schedule this for 30 minutes instead of 45 minutes, and I prefer afternoon meetings after 2pm.\"}]), config=thread_config_5):\n",
    "    # Inspect interrupt object if present\n",
    "    if '__interrupt__' in chunk:\n",
    "        Interrupt_Object = chunk['__interrupt__'][0]\n",
    "        print(\"\\nINTERRUPT OBJECT:\")\n",
    "        print(f\"Action Request: {Interrupt_Object.value[0]['action_request']}\")\n",
    "\n",
    "# Check memory after providing feedback for schedule_meeting\n",
    "print(\"\\nChecking memory after providing feedback for schedule_meeting:\")\n",
    "display_memory_content(store, (\"email_assistant\", \"cal_preferences\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8088757c",
   "metadata": {},
   "source": [
    "Our memory check after providing feedback shows an elegantly simple calendar preference update:\n",
    "\n",
    "```\n",
    "30 minute meetings are preferred, but 15 minute meetings are also acceptable.\n",
    "Afternoon meetings after 2pm are preferred.\n",
    "```\n",
    "\n",
    "The system has:\n",
    "1. Captured both aspects of our feedback (duration and time of day)\n",
    "2. Preserved the existing preference about 15-minute meetings\n",
    "3. Added our preference for afternoon meetings after 2pm as a new line\n",
    "4. Kept the format clean and readable\n",
    "\n",
    "This natural language feedback mechanism creates the same quality of memory updates as direct editing but requires less effort from the user. The system is able to extract structured preferences from unstructured feedback, showing its ability to learn from conversational interactions.\n",
    "\n",
    "Let's accept this revised meeting proposal and move to the email draft:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "545063be",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"\\nSimulating user accepting the {Interrupt_Object.value[0]['action_request']['action']} tool call...\")\n",
    "for chunk in graph.stream(Command(resume=[{\"type\": \"accept\"}]), config=thread_config_5):\n",
    "    # Inspect interrupt object if present\n",
    "    if '__interrupt__' in chunk:\n",
    "        Interrupt_Object = chunk['__interrupt__'][0]\n",
    "        print(\"\\nINTERRUPT OBJECT:\")\n",
    "        print(f\"Action Request: {Interrupt_Object.value[0]['action_request']}\")\n",
    "\n",
    "# Check memory after accepting schedule_meeting after feedback\n",
    "print(\"\\nChecking memory after accepting schedule_meeting after feedback:\")\n",
    "display_memory_content(store, (\"email_assistant\", \"response_preferences\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e72ede94",
   "metadata": {},
   "source": [
    "Now provide feedback for the `write_email` tool call\n",
    "\n",
    "Similar to our meeting feedback, we'll now provide natural language guidance for the email draft:\n",
    "\n",
    "1. We request \"shorter and less formal\" language - a style preference\n",
    "2. We ask for a specific closing statement about looking forward to the meeting\n",
    "3. The system must interpret this guidance and rewrite the email accordingly\n",
    "\n",
    "After providing this feedback, we'll check the response preferences memory to see how these style and structure preferences are captured. We expect to see generalizable guidelines about email brevity, formality, and closing statements added to our preference profile."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9831ad2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"\\nSimulating user providing feedback for the {Interrupt_Object.value[0]['action_request']['action']} tool call...\")\n",
    "for chunk in graph.stream(Command(resume=[{\"type\": \"response\", \"args\": \"Shorter and less formal. Include a closing statement about looking forward to the meeting!\"}]), config=thread_config_5):\n",
    "    # Inspect response_agent most recent message\n",
    "    if 'response_agent' in chunk:\n",
    "        chunk['response_agent']['messages'][-1].pretty_print()\n",
    "    # Inspect interrupt object if present\n",
    "    if '__interrupt__' in chunk:\n",
    "        Interrupt_Object = chunk['__interrupt__'][0]\n",
    "        print(\"\\nINTERRUPT OBJECT:\")\n",
    "        print(f\"Action Request: {Interrupt_Object.value[0]['action_request']}\")\n",
    "\n",
    "# Check memory after providing feedback for write_email\n",
    "print(\"\\nChecking memory after providing feedback for write_email:\")\n",
    "display_memory_content(store, (\"email_assistant\", \"response_preferences\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5b360a2",
   "metadata": {},
   "source": [
    "The memory update after our email feedback shows highly sophisticated learning about both meeting scheduling and email writing preferences:\n",
    "\n",
    "1. The system has added a complete new section to the response preferences entitled \"When writing email responses\" with two key preferences:\n",
    "   - \"Favor shorter and less formal language when possible, unless the context requires formality\"\n",
    "   - \"Include a closing statement expressing that you look forward to the meeting or conversation when confirming appointments\"\n",
    "\n",
    "2. It has also added a new bullet point to the \"When responding to meeting scheduling requests\" section:\n",
    "   - \"When scheduling meetings, prefer afternoon times after 2pm when possible, and default to 30-minute durations unless otherwise specified\"\n",
    "\n",
    "This demonstrates the system's ability to:\n",
    "- Organize learned preferences into appropriate categories\n",
    "- Extract multiple insights from a single feedback instance\n",
    "- Apply meeting preferences to both calendar and email contexts\n",
    "- Capture nuance with appropriate qualifiers (\"when possible,\" \"unless otherwise specified\")\n",
    "- Maintain the hierarchical structure of the memory\n",
    "\n",
    "The resulting email shows all these preferences applied: it's shorter, less formal, includes a closing statement about looking forward to the chat, and correctly references the 30-minute meeting time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c64999e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"\\nSimulating user accepting the {Interrupt_Object.value[0]['action_request']['action']} tool call...\")\n",
    "for chunk in graph.stream(Command(resume=[{\"type\": \"accept\"}]), config=thread_config_5):\n",
    "    # Inspect interrupt object if present\n",
    "    if '__interrupt__' in chunk:\n",
    "        Interrupt_Object = chunk['__interrupt__'][0]\n",
    "        print(\"\\nINTERRUPT OBJECT:\")\n",
    "        print(f\"Action Request: {Interrupt_Object.value[0]['action_request']}\")\n",
    "\n",
    "# Check memory after accepting write_email after feedback\n",
    "print(\"\\nChecking memory after accepting write_email after feedback:\")\n",
    "display_memory_content(store, (\"email_assistant\", \"response_preferences\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f85e63cb",
   "metadata": {},
   "source": [
    "Look at the full message history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa9cf91d",
   "metadata": {},
   "outputs": [],
   "source": [
    "state = graph.get_state(thread_config_5)\n",
    "for m in state.values['messages']:\n",
    "    m.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80ac9df0-cd39-4c32-a073-c2482d9554b6",
   "metadata": {},
   "source": [
    "## Local Deployment\n",
    "\n",
    "You can find this graph with memory integration in the `src/email_assistant` directory:\n",
    "\n",
    "* `src/email_assistant/email_assistant_hitl_memory.py`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a4aa8b0-f8b7-4197-8701-87dda60daa26",
   "metadata": {},
   "source": [
    "Email to test: \n",
    "```\n",
    "{\n",
    "  \"author\": \"Alice Smith <alice.smith@company.com>\",\n",
    "  \"to\": \"John Doe <john.doe@company.com>\",\n",
    "  \"subject\": \"Quick question about API documentation\",\n",
    "  \"email_thread\": \"Hi John,\\nI was reviewing the API documentation for the new authentication service and noticed a few endpoints seem to be missing from the specs. Could you help clarify if this was intentional or if we should update the docs?\\nSpecifically, I'm looking at:\\n- /auth/refresh\\n- /auth/validate\\nThanks!\\nAlice\"\n",
    "}\n",
    "```\n",
    "\n",
    "As before, if you go to [dev.agentinbox.ai](https://dev.agentinbox.ai/), you can easily connect to the graph:\n",
    "\n",
    "   * Graph name: the name from the `langgraph.json` file (`email_assistant_hitl_memory`)\n",
    "   * Graph URL: `http://127.0.0.1:2024/`\n",
    "\n",
    "![inbox](img/agent-inbox-edit.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b075a3ea",
   "metadata": {},
   "source": [
    "The Memory tab in LangGraph Studio offers a real-time view of how your preferences are being captured and updated with each interaction:\n",
    "\n",
    "![studio-img](img/memory-studio.png)\n",
    "\n",
    "Through continued use, the system becomes increasingly personalized:\n",
    "- It learns which emails you want to respond to, be notified about, or ignore\n",
    "- It adapts to your communication style preferences\n",
    "- It remembers your scheduling preferences\n",
    "- It refines its understanding with each interaction\n",
    "\n",
    "This combination of HITL and memory creates a system that balances automation with control - handling routine tasks automatically while learning from your feedback to become more aligned with your preferences over time."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2ad7580",
   "metadata": {},
   "source": [
    "## Hosted Deployment with Gmail Tools\n",
    "\n",
    "If you want to actually run this on your own email, you can deploy the graph with Gmail tools. \n",
    "\n",
    "Set up your Gmail credentials [following here](https://github.com/langchain-ai/agents-from-scratch/blob/main/src/email_assistant/tools/gmail/README.md).\n",
    "\n",
    "There is a graph set up with Gmail tools:\n",
    "\n",
    "```shell\n",
    "python src/email_assistant/email_assistant_hitl_memory_gmail.py\n",
    "```\n",
    "\n",
    "[One of the deployment options is `hosted`](https://langchain-ai.github.io/langgraph/tutorials/deployment/#other-deployment-options), and you can simply connect the deployed graph URL to the Agent Inbox as done with the local deployment.\n",
    "\n",
    "## Improving Memory \n",
    "\n",
    "Our current memory schema and updating is extremely simple: \n",
    "\n",
    "* Our schema is a string\n",
    "* We always overwrite the existing memory with a new string\n",
    " \n",
    "The store can be easily [configured for semantic search](https://langchain-ai.github.io/langgraph/cloud/deployment/semantic_search/) over a collection of memories. \n",
    "\n",
    "Also consider using [LangMem](https://langchain-ai.github.io/langmem/) for more advanced memory management."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "846dbb9b-5c9a-4236-912e-02b3d9f674f6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  },
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
